#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from datetime import datetime


def NormalizeTable1(filepath, index):
    index -= 1
    ht = set()
    infile = open(filepath, "r")
    outfile = open(filepath[:-4] + ".normalized.txt", "w")
    outfile.write("RunID\t" + infile.readline())
    for line in infile:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        if splitline[0] not in ht:
            ht.add(splitline[0])
            index += 1
        run_id = "NODER%08d" % index
        outfile.write(
            run_id + "\t" + splitline[0] + "\txushua_" + run_id + "\t" + splitline[2] + "\t" + splitline[3] + "\t" +
            splitline[4] + "\n")
        outfile.flush()
    outfile.close()
    infile.close()


def NormalizeTable2(filepath):
    ht = {}
    infile1 = open(filepath, "r")
    infile2 = open(filepath[:-5] + "1.normalized.txt", "r")
    outfile = open(filepath[:-4] + ".normalized.txt", "w")

    infile2.readline()
    for line in infile2:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        ht[splitline[1]] = splitline[0]
    infile2.close()

    outfile.write("RunID\t" + infile1.readline())
    for line in infile1:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        outfile.write(ht[splitline[1]] + "\t" + line + "\n")
        outfile.flush()
    outfile.close()
    infile1.close()


def SearchSample(filepath):
    client = MongoClient("mongo-node", 27017)
    collection = client["node"]["sample"]

    infile = open(filepath, "r")
    infile.readline()
    for line in infile:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        result = collection.count({"name": splitline[5]})
        print(splitline[5] + "\t" + str(result))
    infile.close()
    client.close()


def GenerateSampleFile(filepathlist, index):
    ht = {}
    outfile = open("C:/Users/Genesis/Desktop/Xushuhua/samplelist.txt", "w")
    for filepath in filepathlist:
        infile = open(filepath, "r")
        outfile1 = open(filepath[:-4] + "2.txt", "w")
        outfile1.write(infile.readline().rstrip("\n") + "\tSampleID\n")
        for line in infile:
            line = line.rstrip("\n")
            splitline = line.split("\t")
            if splitline[5] not in ht:
                index += 1
                sample_id = "NODES%08d" % index
                ht[splitline[5]] = sample_id
                outfile.write(sample_id + "\t" + splitline[5] + "\t" + splitline[3] + "\n")
                outfile.flush()
            outfile1.write(line + "\t" + ht[splitline[5]] + "\n")
            outfile1.flush()
        infile.close()
        outfile1.close()
    outfile.close()


def InsertSamples(filepath):
    client = MongoClient("mongo-node", 27017)
    collection = client["node"]["sample"]

    infile = open(filepath, "r")
    for line in infile:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        insert_ht = {}
        insert_ht["_class"] = "org.scbit.lsbi.node.pojo.Sample"
        insert_ht["sap_no"] = splitline[0]
        insert_ht["name"] = splitline[1]
        insert_ht["description"] = splitline[2]
        insert_ht["taxid"] = "2"
        insert_ht["taxonomy"] = "Bacteria"
        insert_ht["organism"] = "Bacteria"
        insert_ht["division"] = "Bacteria"
        insert_ht["attributes"] = {"Source": "Gut"}
        insert_ht["provider"] = None
        insert_ht["related_links"] = []
        insert_ht["creator"] = "TEA253TXEBCEXNKXGFZHOV3GAA"
        insert_ht["createDate"] = datetime.strptime("2017-08-22T21:00:00Z", "%Y-%m-%dT%H:%M:%SZ")
        insert_ht["updater"] = "TEA253TXEBCEXNKXGFZHOV3GAA"
        insert_ht["update_date"] = datetime.strptime("2017-08-22T21:00:00Z", "%Y-%m-%dT%H:%M:%SZ")
        insert_ht["security"] = "Public"
        insert_ht["public_date"] = None
        insert_ht["hit_num"] = 0
        insert_ht["operator"] = "RTYAFSWTP5GI5EP3BUU5YSAQIU"
        insert_ht["submitter"] = {"firstName": "Shuhua", "lastName": "Xu", "orgName": "PICB",
                                  "deptName": "Shuhua Xu Lab", "email": ""}
        collection.insert_one(insert_ht)
    infile.close()
    client.close()


def InsertRuns(filepathlist):
    client = MongoClient("mongo-node", 27017)
    collection = client["node"]["run"]
    for filepath in filepathlist:
        infile = open(filepath, "r")
        infile.readline()
        for line in infile:
            line = line.rstrip("\n")
            splitline = line.split("\t")
            insert_ht = {}
            insert_ht["_class"] = "org.scbit.lsbi.node.pojo.Run"
            insert_ht["run_no"] = splitline[0]
            insert_ht["exp_no"] = splitline[4]
            insert_ht["sap_no"] = splitline[6]
            insert_ht["name"] = splitline[2]
            insert_ht["description"] = splitline[3]
            insert_ht["related_links"] = []
            insert_ht["creator"] = "TEA253TXEBCEXNKXGFZHOV3GAA"
            insert_ht["createDate"] = datetime.strptime("2017-08-22T21:00:00Z", "%Y-%m-%dT%H:%M:%SZ")
            insert_ht["updater"] = "TEA253TXEBCEXNKXGFZHOV3GAA"
            insert_ht["update_date"] = datetime.strptime("2017-08-22T21:00:00Z", "%Y-%m-%dT%H:%M:%SZ")
            insert_ht["security"] = "Restricted"
            insert_ht["public_date"] = None
            insert_ht["hit_num"] = 0
            insert_ht["operator"] = "RTYAFSWTP5GI5EP3BUU5YSAQIU"
            insert_ht["submitter"] = {"firstName": "Shuhua", "lastName": "Xu", "orgName": "PICB",
                                      "deptName": "Shuhua Xu Lab", "email": ""}
            collection.insert_one(insert_ht)
        infile.close()
    client.close()


def UpdateData(filepathlist):
    client = MongoClient("mongo-node", 27017)
    collection = client["node"]["data"]
    for filepath in filepathlist:
        infile = open(filepath, "r")
        infile.readline()
        for line in infile:
            line = line.rstrip("\n")
            splitline = line.split("\t")
            collection.update_one({"dat_no": splitline[1]}, {
                "$set": {"run_no": splitline[0], "archived": "yes", "updater": "TEA253TXEBCEXNKXGFZHOV3GAA",
                         "update_date": datetime.strptime("2017-08-22T21:00:00Z", "%Y-%m-%dT%H:%M:%SZ"),
                         "security": "Restricted", "public_date": None, "operator": "RTYAFSWTP5GI5EP3BUU5YSAQIU"}})
        infile.close()
    client.close()


def UpdateRun(filepathlist):
    client = MongoClient("mongo-node", 27017)
    collection = client["node"]["run"]
    for filepath in filepathlist:
        infile = open(filepath, "r")
        infile.readline()
        for line in infile:
            line = line.rstrip("\n")
            splitline = line.split("\t")
            collection.update_one({"run_no": splitline[0]}, {"$set": {"sap_no": splitline[6]}})
        infile.close()
    client.close()


if __name__ == "__main__":
    # NormalizeTable1("C:/Users/Genesis/Desktop/Xushuhua/NODEP00000052.table1.txt", 714)
    # NormalizeTable1("C:/Users/Genesis/Desktop/Xushuhua/NODEP00000053.table1.txt", 906)
    # NormalizeTable2("C:/Users/Genesis/Desktop/Xushuhua/NODEP00000052.table2.txt")
    # NormalizeTable2("C:/Users/Genesis/Desktop/Xushuhua/NODEP00000053.table2.txt")
    # SearchSample("NODEP00000052.table1.normalized.txt")
    # SearchSample("NODEP00000053.table1.normalized.txt")
    # GenerateSampleFile(["C:/Users/Genesis/Desktop/Xushuhua/NODEP00000052.table1.normalized.txt",
    #                     "C:/Users/Genesis/Desktop/Xushuhua/NODEP00000053.table1.normalized.txt"], 803)

    # InsertSamples("samplelist.txt")
    # InsertRuns(["NODEP00000052.table1.normalized2.txt", "NODEP00000053.table1.normalized2.txt"])
    # UpdateData(["NODEP00000052.table2.normalized.txt", "NODEP00000053.table2.normalized.txt"])
    UpdateRun(["NODEP00000052.table1.normalized2.txt", "NODEP00000053.table1.normalized2.txt"])
